from collections.abc import Iterable, Iterator
from typing import TypeVarTuple

from zkl_aiutils_datasets.basics import Dataset, DatasetIterator, pause, resume, skip
from zkl_aiutils_datasets.wrapping import wrap_dataset, wrap_dataset_iterator

AnySamples = TypeVarTuple('AnySamples')


class ZippedIterator(DatasetIterator[tuple[*AnySamples]]):
    def __init__(self, iterators: Iterable[Iterator | Iterable]):
        self._iterators = tuple(wrap_dataset_iterator(iterator) for iterator in iterators)

    def __next__(self) -> tuple[*AnySamples]:
        samples = []
        for iterator in self._iterators:
            samples.append(next(iterator))  # may raise StopIteration
        # noinspection PyTypeChecker
        return tuple(samples)

    def __skip__(self, samples_n: int):
        for iterator in self._iterators:
            skip(iterator, samples_n)

    def __pause__(self) -> tuple:
        return tuple(pause(iterator) for iterator in self._iterators)


class ZippedDataset(Dataset[tuple[*AnySamples]]):
    def __init__(self, datasets: Iterable[Iterable]):
        self._datasets = tuple(wrap_dataset(dataset) for dataset in datasets)

    def __len__(self) -> int:
        return min(len(dataset) for dataset in self._datasets)

    def __getitem__(self, index: int, /) -> tuple[*AnySamples]:
        # noinspection PyTypeChecker
        return tuple(dataset[index] for dataset in self._datasets)

    def __iter__(self) -> ZippedIterator[*AnySamples]:
        return ZippedIterator(iter(dataset) for dataset in self._datasets)

    def __resume__(self, state: tuple | None = None, /) -> ZippedIterator[*AnySamples]:
        return ZippedIterator(resume(dataset, state) for dataset, state in zip(self._datasets, state))
